"""
支持滑动窗口的数据加载和预处理模块
"""
import os
import numpy as np
import scipy.io
import torch
from config import Config
from torch.utils.data import Dataset, DataLoader as TorchDataLoader

class SignalDataset(Dataset):
    """信号数据集类"""
    
    def __init__(self, config, start_idx, end_idx, sliding_step=1):
        self.config = config
        self.start_idx = start_idx
        self.end_idx = end_idx
        self.sliding_step = sliding_step
        self.seq_len = config.SEQ_LEN
        self.file_prefix = config.FILE_PREFIX
        self.file_extension = config.FILE_EXTENSION
        self.data_folder = config.DATA_FOLDER
        
        # 预加载所有文件路径
        self.file_paths = []
        for i in range(start_idx, end_idx + 1):
            file_path = os.path.join(
                self.data_folder,
                f"{self.file_prefix}{i}{self.file_extension}"
            )
            if os.path.exists(file_path):
                self.file_paths.append(file_path)
    
    def __len__(self):
        return len(self.file_paths) - self.seq_len + 1
    
    def __getitem__(self, idx):
        # 计算起始文件索引
        start_file_idx = idx * self.sliding_step
        
        # 加载序列数据
        result_matrix = np.zeros((8, 8, self.seq_len), dtype=np.complex64)
        
        for t in range(self.seq_len):
            file_path = self.file_paths[start_file_idx + t]
            try:
                mat_data = scipy.io.loadmat(file_path)
                if "p" in mat_data:
                    Y = mat_data["p"]
                    if Y.size == 64:
                        Y_reshaped = Y.reshape(8, 8)
                        result_matrix[:, :, t] = Y_reshaped
            except Exception as e:
                print(f"错误: 读取文件 {file_path} 失败: {e}")
        
        # 处理复数数据
        Y_real = np.real(result_matrix).reshape(64, self.seq_len)
        Y_imag = np.imag(result_matrix).reshape(64, self.seq_len)
        Y_combined = np.concatenate((Y_real, Y_imag), axis=0)
        
        return torch.from_numpy(Y_combined).float()

class DataLoader:
    """支持滑动窗口的数据加载器类"""
    
    def __init__(self, config=None, sliding_step=5):
        self.config = config or Config()
        self.sliding_step = sliding_step
        self.endmember_matrix = self._load_endmember_matrix()
        
        # 创建数据集
        self.train_dataset = SignalDataset(
            self.config,
            start_idx=1,
            end_idx=self.config.MAX_TRAIN_BATCHES * self.config.BATCH_SIZE,
            sliding_step=self.sliding_step
        )
        
        self.test_dataset = SignalDataset(
            self.config,
            start_idx=self.config.MAX_TRAIN_BATCHES * self.config.BATCH_SIZE + 1,
            end_idx=self.config.TOTAL_DATA,
            sliding_step=self.sliding_step
        )
        
        # 创建数据加载器
        self.train_loader = TorchDataLoader(
            self.train_dataset,
            batch_size=self.config.BATCH_SIZE,
            shuffle=True,
            num_workers=4,
            pin_memory=True
        )
        
        self.test_loader = TorchDataLoader(
            self.test_dataset,
            batch_size=self.config.BATCH_SIZE,
            shuffle=False,
            num_workers=4,
            pin_memory=True
        )
        
        print(f"训练集大小: {len(self.train_dataset)}")
        print(f"测试集大小: {len(self.test_dataset)}")
    
    def _load_endmember_matrix(self):
        """加载端元矩阵D"""
        try:
            mat_data = scipy.io.loadmat(self.config.ENDMEMBER_PATH)
            D = mat_data["phasecha"]
            
            if D.shape[0] == 64:
                D_reshaped = D.reshape(64, self.config.K)
            else:
                raise ValueError(f"端元矩阵的行数 {D.shape[0]} 与预期的64不匹配")
            
            # 分离实部和虚部
            D_real = np.real(D_reshaped)
            D_imag = np.imag(D_reshaped)
            D_combined = np.concatenate((D_real, D_imag), axis=0)
            
            return torch.from_numpy(D_combined).float().to(self.config.DEVICE)
            
        except Exception as e:
            raise RuntimeError(f"加载端元矩阵失败: {e}")
    
    def get_train_loader(self):
        """获取训练数据加载器"""
        return self.train_loader
    
    def get_test_loader(self):
        """获取测试数据加载器"""
        return self.test_loader
    
    def get_endmember_matrix(self):
        """获取端元矩阵"""
        return self.endmember_matrix
    
    def _count_available_files(self):
        """计算可用的文件数量"""
        file_count = 0
        for i in range(1, self.config.TOTAL_DATA + 1):
            file_path = os.path.join(
                self.config.DATA_FOLDER,
                f"{self.config.FILE_PREFIX}{i}{self.config.FILE_EXTENSION}"
            )
            if os.path.exists(file_path):
                file_count += 1
        return file_count
    
    def _calculate_sliding_sequences(self):
        """计算滑动窗口可以产生的序列数量"""
        # 最后一个序列的起始位置不能超过 available_files - seq_len + 1
        max_start_idx = self._count_available_files() - self.config.SEQ_LEN + 1
        # 滑动窗口序列数量 = (max_start_idx - 1) // sliding_step + 1
        return max_start_idx // self.sliding_step if max_start_idx > 0 else 0
    
    def print_sliding_info(self):
        """打印滑动窗口信息"""
        print(f"\n=== 滑动窗口数据加载信息 ===")
        print(f"滑动步长: {self.sliding_step}")
        print(f"序列长度: {self.config.SEQ_LEN}")
        print(f"批次大小: {self.config.BATCH_SIZE}")
        print(f"可用文件数: {self._count_available_files()}")
        print(f"滑动窗口序列数: {self._calculate_sliding_sequences()}")
        print(f"最大批次数: {self._calculate_sliding_sequences() // self.config.BATCH_SIZE}")
        print(f"数据增强倍数: {self._calculate_sliding_sequences() / (self._count_available_files() // self.config.SEQ_LEN):.1f}x")
        
        # 显示前几个序列的文件范围
        print(f"\n前几个序列的文件范围:")
        for i in range(min(5, self._calculate_sliding_sequences())):
            start_idx = i * self.sliding_step + 1
            end_idx = start_idx + self.config.SEQ_LEN - 1
            print(f"  序列 {i}: 文件 {start_idx}-{end_idx}") 